package com.allin.mall.gateway.filter;

import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;

@Component
@ConditionalOnProperty(value = "mall.response.encrypt.enable", havingValue = "true")
@Order(-1)
public class ResponseEncryptFilter implements GlobalFilter {
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        ServerHttpResponse originalResponse = exchange.getResponse();
        originalResponse.getHeaders().setContentType(MediaType.APPLICATION_JSON);

        ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(originalResponse) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                    return super.writeWith(fluxBody.collectList().flatMap(dataBuffer -> {
                       // 合并所有缓冲区
                        int totalSize = dataBuffer.stream().mapToInt(DataBuffer::readableByteCount).sum();
                        byte[] combined = new byte[totalSize];
                        int offset = 0;

                        for (DataBuffer buffer : dataBuffer) {
                            int length = buffer.readableByteCount();
                            buffer.read(combined, offset, length);
                            offset += length;
                            DataBufferUtils.release(buffer);
                        }

                        // 加密内容
                        String aOriginalResponse = new String(combined, StandardCharsets.UTF_8);
                        String encryptedResponse = String.format("{\"e\":\"%s\"}", doEncrypt(aOriginalResponse));

                        byte[] encryptedBytes = encryptedResponse.getBytes(StandardCharsets.UTF_8);
                        originalResponse.getHeaders().setContentLength(encryptedBytes.length);
                        // 返回加密后的内容
                        return Mono.just(bufferFactory().wrap(encryptedBytes));
                    }));
                }
                return super.writeWith(body);
            }
        };
        return chain.filter(exchange.mutate().response(decoratedResponse).build());
    }

    /**
     * 对响应体进行加密
     * @param original
     * @return
     */
    private String doEncrypt(String original) {
        // TODO: ...
        return original;
    }
}
